# -*- coding: utf-8 -*-

# from functools import partial

# import matplotlib.pyplot as plt
import numpy as np
import segyio

def get_index(arr, val):
    """
    查找数组中值为 val 的元素的下标，如果找不到，则返回离其最近的上一个值
    :param arr: 深度数组
    :param val: 要查找的数据
    :return: 下标
    """
    indices = np.where(arr == val)[0]
    if indices.size > 0:
        # 如果找到了 deep，获取其第一个下标
        index = indices[0]
    else:
        # 如果没找到 deep，向下取最近的数据的下标
        diff = arr - val
        valid_indices = np.where(diff < 0)[0]
        if valid_indices.size > 0:
            index = valid_indices[diff[valid_indices].argmax()]
        else:
            # 如果所有元素都大于 val，取第一个元素的下标
            index = 0

    return index

class Seismic(object):

    def __init__(self, data, dtype=float, params=None):

        if params is None:
            params = {}

        self.params = params
        self.data = np.asarray(data, dtype=dtype)
        self.ntraces = params.get('ntraces', self.data.shape[-2])
        self.dt = params.get('dt', 0.001)
        self.cdp = params.get('cdp')
        self.fb = params.get('fb')
        self.nsamples = int(self.data.shape[-1])
        self.data = np.squeeze(self.data)
        self.cdpx = params.get('cdpx')
        self.cdpy = params.get('cdpy')
        self.wellx = params.get('wellx')
        self.welly = params.get('welly')
        self.inlin = params.get('inlin')
        self.crosslin = params.get('crosslin')
        self.hcf = params.get('hcf')
        # 修改时间长度
        self.tbasis = np.arange(0, self.nsamples * self.dt, self.dt)
        return

    @classmethod
    def from_segyio(cls, segy_file, params=None):
        params = {}
        with segyio.open(segy_file, strict=False, ignore_geometry=True) as s:
            # Read the data.
            data = np.stack([t.astype(float) for t in s.trace])

            # Get the trace numbers.
            cdp = np.array([t[segyio.TraceField.CDP_X] for t in s.header])
            fb = np.array(
                [t[segyio.TraceField.MuteTimeStart] for t in s.header])
            cdpx = np.array([t[segyio.TraceField.CDP_X] for t in s.header])
            cdpy = np.array([t[segyio.TraceField.CDP_Y] for t in s.header])
            wellx = np.array([t[segyio.TraceField.GroupX] for t in s.header])
            welly = np.array([t[segyio.TraceField.GroupY] for t in s.header])
            inlin = np.array(
                [t[segyio.TraceField.INLINE_3D] for t in s.header])
            crosslin = np.array(
                [t[segyio.TraceField.CROSSLINE_3D] for t in s.header])
            hcf = np.array(
                [t[segyio.TraceField.HighCutFrequency] for t in s.header])

            params['cdpx'] = cdpx
            params['cdpy'] = cdpy
            params['wellx'] = wellx
            params['welly'] = welly
            params['cdp'] = cdp
            params['fb'] = fb
            params['inlin'] = inlin
            params['crosslin'] = crosslin
            params['hcf'] = hcf

            params['dt'] = s.bin[segyio.BinField.Interval] / 1000

        return cls(data, dtype=float, params=params)

    @classmethod
    def from_segy_with_segyio(cls, segy_file, params=None):
        return cls.from_segyio(segy_file, params=params)

    from_segy = from_segy_with_segyio

    def get_data(self,
                 l=1):

        return self.data

    def wiggle_plot(self,
                    l=1,
                    ax=None,
                    skip=1,
                    skip2=2,
                    perc=99.0,
                    gain=1.0,
                    rgb=(0, 0, 0),
                    alpha=0.5,
                    lw=0.2,
                    ts=0.0,
                    td=6.0,
                    kb=0.0,
                    dir=0,
                    stkw=100,
                    stkdep=0,
                    fcx=1,
                    dep=[],
                    zu=[],
                    font=[],
                    font1=[],
                    fzt=[],
                    py=[],
                    ywdh=[],
                    ywdhpy=[],
                    lped=0.0,
                    lpsd=0.0,
                    ):

        data = self.get_data(l)
        rgba = list(rgb) + [alpha]
        sc = np.percentile(data, perc)  # Normalization factor，计算数据集中特定百分位数的值
        stp = int(ts * 1000.0 / self.dt)
        etp = int(td * 1000.0 / self.dt)
        wigdata = data[::skip, stp: etp: skip2]
        brfk = 1
        self.cdp = self.cdp + kb
        xpos = np.array(self.cdp)
        brfp = 0
        fbmax = 0.0
        fbmin = 8000.0
        for brfjj, x, trace in zip(range(0, len(xpos)), xpos, wigdata):
            # Compute high resolution trace.
            if x > lped + kb or x < lpsd + kb:
                brfp = brfp + 1
                brfk = brfk + 1
                print('跳过第', brfp, '道，共', len(xpos), '道')
                continue
            if self.fb[brfjj] > fbmax:
                fbmax = self.fb[brfjj]
            if self.fb[brfjj] < fbmin:
                fbmin = self.fb[brfjj]
            if brfk == xpos.size:
                brfk = brfk - 1
            if dir == 0:
                amp = x - gain * (xpos[brfk] - xpos[brfk - 1]) * trace / sc
            else:

                amp = x - gain * (xpos[brfk - 1] - xpos[brfk]) * trace / sc

            t = np.arange(ts * 1000, td * 1000, self.dt * skip2)
            hypertime = np.linspace(t[0], t[-1], (10 * t.size - 1) + 1)
            hyperamp = np.interp(hypertime, t, amp)

            # Plot the line, then the fill.
            for brfi in range(len(hyperamp)):
                if hypertime[brfi] < self.fb[brfp]:
                    hyperamp[brfi] = x
                if x < stkdep:
                    if hypertime[brfi] > self.fb[brfp] + stkw:
                        hyperamp[brfi] = x
            ax.plot(hyperamp, hypertime, 'k', lw=1.2 * lw)
            ax.fill_betweenx(hypertime, hyperamp, x,
                             where=hyperamp < x,
                             facecolor=rgba,
                             lw=0,
                             )

            # print('{:.2f}'.format(brfp / len(xpos) * 100), '%')

            brfk = brfk + 1
            brfp = brfp + 1
            # print('绘制第', brfp, '道，共', len(xpos), '道')

        ax.plot(self.cdp, self.fb, 'r', lw=2.5 * lw)
        # 绘制蓝线
        index = get_index(self.cdp, stkdep)
        line_cdp = self.cdp[:index + 1]
        line_fb = self.fb[:index + 1] + stkw
        ax.plot(line_cdp, line_fb, 'b', lw=2.5 * lw)

        index2 = get_index(hypertime, self.fb[index] + stkw)
        line_time = hypertime[index2:]
        line_deep = np.full(line_time.size, self.cdp[index])
        ax.plot(line_deep, line_time, 'b', lw=2.5 * lw)

        ax.invert_yaxis()

        brfpp = 0
        fctime = []
        if fcx == 1:
            fctime = np.arange(len(dep))
            for ii in dep:
                # if ii > lped + kb or ii < lpsd + kb:
                #     brfpp = brfpp + 1
                #     continue
                for jj in range(1, len(self.cdp)):
                    if ii == self.cdp[jj]:
                        fctime[brfpp] = self.fb[jj]
                        brfpp = brfpp + 1
                    elif ((ii > self.cdp[jj - 1] and ii < self.cdp[jj]) or (
                            ii < self.cdp[jj - 1] and ii > self.cdp[jj])):
                        fctime[brfpp] = (ii - self.cdp[jj]) * (
                            self.fb[jj - 1] - self.fb[jj]) / (
                            (self.cdp[jj - 1] - self.cdp[jj])) + self.fb[jj]
                        brfpp = brfpp + 1

            for ii in range(0, len(dep)):
                if dep[ii] > lped + kb or dep[ii] < lpsd + kb:
                    continue
                ax.plot([dep[ii], dep[ii], 0], [0, fctime[ii],
                                                fctime[ii]], 'r', lw=2.5 * lw)
                ax.text(ywdhpy[ii], fctime[ii], ywdh[ii],
                        transform=ax.transData, verticalalignment='center',
                        horizontalalignment='center',
                        fontsize=font1[ii],
                        rotation=270, color='k')
                if ii >= 1 and (ii < (len(dep) - 1)):
                    if (dep[ii] - dep[ii - 1] < font[ii] * 16) and (
                            (dep[ii + 1] - dep[ii] < font[ii] * 16)):
                        ax.text(dep[ii] + font[ii] * 4 + py[ii],
                                fzt[ii], zu[ii], transform=ax.transData,
                                fontsize=font[ii],
                                rotation=270, color='r')
                    elif (dep[ii] - dep[ii - 1] < font[ii] * 16) and (
                            (dep[ii + 1] - dep[ii] > font[ii] * 16)):
                        ax.text(dep[ii] + font[ii] * 8 + py[ii],
                                + fzt[ii], zu[ii], transform=ax.transData,
                                fontsize=font[ii],
                                rotation=270, color='r')
                    elif (dep[ii] - dep[ii - 1] > font[ii] * 16) and (
                            (dep[ii + 1] - dep[ii] < font[ii] * 16)):
                        ax.text(dep[ii],  fzt[ii] + py[ii], zu[ii],
                                transform=ax.transData,
                                fontsize=font[ii],
                                rotation=270, color='r')
                    elif (dep[ii] - dep[ii - 1] > font[ii] * 16) and (
                            (dep[ii + 1] - dep[ii] > font[ii] * 16)):
                        ax.text(dep[ii] + font[ii] * 8 + py[ii],  fzt[ii],
                                zu[ii], transform=ax.transData,
                                fontsize=font[ii],
                                rotation=270, color='r')
                else:
                    ax.text(dep[ii] + py[ii],   fzt[ii], zu[ii],
                            transform=ax.transData, fontsize=font[ii],
                            rotation=270, color='r')
        # if dir == 0:
        ax.set_xlim(lpsd + kb - 20, lped + kb + 20)
        # else:
        #     ax.set_xlim(lpsd-20, xpos[0] + 20)
        ax.invert_xaxis()
        return [fbmin, fbmax, ax, fctime]

    def wiggle_plot1(self,
                     l=1,
                     ax=None,
                     skip=1,
                     skip2=2,
                     perc=99.0,
                     gain=1.0,
                     rgb=(0, 0, 0),
                     alpha=0.5,
                     lw=0.2,
                     ts=0.0,
                     td=6.0,
                     times=0,
                     fcx=1,
                     fctime=[],
                     fbmin=0,
                     fbmax=0,
                     ):

        data = self.get_data(l)
        rgba = list(rgb) + [alpha]
        sc = np.percentile(data, perc)  # Normalization factor
        stp = int((ts * 1000.0 + times) / self.dt)
        etp = int((td * 1000.0 + times) / self.dt)
        wigdata = data[::skip, stp: etp: skip2]

        xpos = np.arange(20, 220, 20)

        for x, trace in zip(xpos, wigdata):
            # Compute high resolution trace.
            amp = x - gain * 20 * trace / sc

            t = np.arange(ts * 1000, td * 1000, self.dt * skip2)
            hypertime = np.linspace(t[0], t[-1], (10 * t.size - 1) + 1)
            hyperamp = np.interp(hypertime, t, amp)

            ax.plot(hyperamp, hypertime, 'k', lw=1.2 * lw)
            ax.fill_betweenx(hypertime, hyperamp, x,
                             where=hyperamp < x,
                             facecolor=rgba,
                             lw=0,
                             )
        if fcx == 1:
            for ii in range(0, len(fctime)):
                if fctime[ii] > fbmax or fctime[ii] < fbmin:
                    continue
                ax.plot([0, 220], [fctime[ii], fctime[ii]], 'r', lw=2.5 * lw)
        ax.set_xlim(0, 220)
        ax.invert_yaxis()
        ax.invert_xaxis()

        return ax

    def wiggle_plot2(self,
                     l=1,
                     ax=None,
                     skip=1,
                     skip2=2,
                     perc=99.0,
                     gain=1.0,
                     rgb=(0, 0, 0),
                     alpha=0.5,
                     lw=0.2,
                     ts=0.0,
                     td=6.0,
                     times=0,
                     stt=0,
                     sein=0,
                     fcx=1,
                     fctime=[],
                     fbmin=0,
                     fbmax=0,
                     ):

        data = self.get_data(l)
        rgba = list(rgb) + [alpha]
        sc = np.percentile(data, perc)  # Normalization factor
        stp = int((ts * 1000.0 + times) / self.dt)
        etp = int((td * 1000.0 + times) / self.dt)

        wigdata = data[stt:stt + sein:skip, stp: etp: skip2]

        xpos = np.arange(20, 20 * sein + 20, 20)

        for x, trace in zip(xpos, wigdata):
            # Compute high resolution trace.
            amp = x + gain * 20 * trace / sc

            t = np.arange(stp * self.dt, etp * self.dt, self.dt * skip2)
            hypertime = np.linspace(t[0], t[-1], (10 * t.size - 1) + 1)
            hyperamp = np.interp(hypertime, t, amp)

            ax.plot(hyperamp, hypertime, 'k', lw=1.2 * lw)
            ax.fill_betweenx(hypertime, hyperamp, x,
                             where=hyperamp > x,
                             facecolor=rgba,
                             lw=0,
                             )
        if fcx == 1:
            for ii in range(0, len(fctime)):
                if fctime[ii] > fbmax or fctime[ii] < fbmin:
                    continue
                ax.plot([0, 100], [fctime[ii] + times,
                        fctime[ii] + times], 'r', lw=2.5 * lw)

        # ax.plot([0,0],[0,fctime[ii] + times],'r', lw=3.0 * lw)
        ax.set_xlim(0, 20 * sein + 20)
        ax.invert_yaxis()
        # ax.invert_xaxis()

        return ax
